#include "MultigridSolver.h"
#include <chrono>
#include <iomanip>
#include <iostream>
#include <fstream>
#include <nlohmann/json.hpp>
#include <Eigen/Dense>
#include <cmath>
#include <Eigen/SparseLU>

using json = nlohmann::json;
using namespace Eigen;

MultigridSolver::MultigridSolver(Grid& grid, BoundaryConditions& bc, const std::string& filename)
    : grid(grid), bc(bc) {
    n = grid.getN();
    loadFromJson(filename);
    getExact_solution();
    /*std::cout << "n: " << n << std::endl;*/
}

void MultigridSolver::loadFromJson(const std::string& filename) {
    std::ifstream inputFile(filename);
    json config;
    inputFile >> config;

    dimension = config["grid"]["dimension"];
    if (dimension == 1) {
        std::string initial_guess_str = config["initial_guess"];
        if(initial_guess_str == "zero"){
            initial_guess.resize(n+1, 0.0);
        }
    } else if (dimension == 2) {
        std::string initial_guess_str = config["initial_guess"];
        if(initial_guess_str == "zero"){
            initial_guess.resize((n+1)*(n+1), 0.0);
        }
    }

    method = config["solver"]["method"];
    tolerance = config["solver"]["tolerance"];
    maxIterations = config["solver"]["max_iter"];
    restriction = config["solver"]["restriction"];
    interpolation = config["solver"]["interpolation"];
    v1 = config["solver"]["v1"];
    v2 = config["solver"]["v2"];
    if(bc.getType()=="mixed"){
        types = config["boundary_conditions"].at("types").get<std::vector<std::string>>();
    }
}

VectorXd MultigridSolver::getExact_solution() {
    if (dimension == 1) {
        std::function<double(double)> exact_fun = grid.getExactSolution_1D();
        auto g = grid.getGrid_1D();  // 网格点
        int n = g.size();
        // 计算真实解
        exact_solution = VectorXd(n);
        for (int i = 0; i < n; ++i) {
            exact_solution(i) = exact_fun(g[i]);
        }
        /*// 打印1D真实解
        std::cout << "Exact solution (1D):\n" << exact_solution.transpose() << "\n";*/
    }
    else if (dimension == 2) {
        std::function<double(double, double)> exact_fun = grid.getExactSolution_2D();
        auto g = grid.getGrid_2D();  // 网格点
        int n = g.size();
        // 计算真实解
        exact_solution = VectorXd(n);
        for (int i = 0; i < n; ++i) {
            exact_solution(i) = exact_fun(g[i][0], g[i][1]);
        }
        /*// 打印2D真实解
        std::cout << "Exact solution (2D):\n" << exact_solution.transpose() << "\n";*/
    }
    
    return exact_solution;
}


VectorXd MultigridSolver::relax(int numIterations, Eigen::SparseMatrix<double> A, VectorXd v, VectorXd f) {
    double omega = 2.0 / 3.0;

    for (int k = 0; k < numIterations; ++k) {
        VectorXd r = f - A * v;

        for (int i = 0; i < v.size(); ++i) {
            if (A.coeff(i, i) != 0) { // 防止除零错误
                v(i) += omega * r(i) / A.coeff(i, i);
            }
        }
    }

    return v;
}

VectorXd MultigridSolver::restrict(VectorXd v) {
    VectorXd vCoarse;
    if(dimension==1){
        int nCoarse = (v.size() + 1) / 2;
        vCoarse.resize(nCoarse);

        if (restriction == "full weighting") {
            vCoarse(0) = v[0];
            for (int i = 1; i < nCoarse - 1; ++i) {
            vCoarse(i) = (v(2 * i - 1) + 2 * v(2 * i) + v(2 * i + 1)) / 4.0;
            }
            vCoarse(nCoarse - 1) = v[2*nCoarse-2];
        } else if (restriction == "injection") {
            for (int i = 0; i < nCoarse; ++i) {
                vCoarse(i) = v[2 * i];
            }
        }
    }
    else if(dimension==2){
        int n1 = static_cast<int>(std::sqrt(static_cast<double>(v.size())));
        int n2 = static_cast<int>((n1 + 1) / 2);
        int nCoarse = static_cast<int>(std::pow(n2, 2));
        vCoarse.resize(nCoarse);

        if (restriction == "full weighting") {
            vCoarse(0) = v[0];
            for (int i = 1; i < nCoarse - 1; ++i) {
            vCoarse(i) = (v(2 * i - 1) + 2 * v(2 * i) + v(2 * i + 1)) / 4.0;
            }
            vCoarse(nCoarse - 1) = v[2*nCoarse];
        } else if (restriction == "injection") {
            for (int i = 0; i < nCoarse; ++i) {
                vCoarse(i) = v[2 * i + (i / n2) * (n1 - 1)];
            }
        }
    }
    
    return vCoarse;
}

VectorXd MultigridSolver::interpolate(VectorXd v) {
    VectorXd vFine;

    if(dimension==1){
        int nFine = v.size() * 2 - 1;
        vFine.resize(nFine);
        if (interpolation == "linear") {
            for (int i = 0; i < v.size(); ++i) {
                vFine(2 * i) = v(i);
                if (2 * i + 1 < nFine) {
                    vFine(2 * i + 1) = 0.5 * (v(i) + v(i + 1));
                }
            }
        } else if (interpolation == "quadratic") {
            vFine(1) = 3.0/8.0 * v(0) + 3.0/4.0 * v(1) - 1.0/8.0 * v(2);
            vFine(nFine - 2) = -1.0/8.0 * v(v.size() - 3) + 3.0/4.0 * v(v.size() - 2) + 3.0/8.0 * v(v.size() - 1);
        
            for (int i = 0; i < v.size(); ++i) {
                vFine(2 * i) = v(i);
        
                if (1 < 2 * i + 1 && 2 * i + 1 < nFine - 2) {
                    vFine(2 * i + 1) = 
                        -1.0/16.0 * v(i - 1) + 
                        9.0/16.0 * v(i) + 
                        9.0/16.0 * v(i + 1) - 
                        1.0/16.0 * v(i + 2);
                }
            }
        }    
    }
    else if(dimension==2){
        int n1 = static_cast<int>(std::sqrt(static_cast<double>(v.size())));
        int n2 = static_cast<int>(n1 * 2 - 1);
        int nFine = static_cast<int>(std::pow(n2, 2));
        vFine.resize(nFine);
        if (interpolation == "linear") {
            for (int i = 0; i < v.size(); ++i) {
                vFine(2 * i) = v(i);
                if (2 * i + 1 < nFine) {
                    vFine(2 * i + 1) = 0.5 * (v(i) + v(i + 1));
                }
            }
        } else if (interpolation == "quadratic") {
            vFine(1) = 3.0/8.0 * v(0) + 3.0/4.0 * v(1) - 1.0/8.0 * v(2);
            vFine(nFine - 2) = -1.0/8.0 * v(v.size() - 3) + 3.0/4.0 * v(v.size() - 2) + 3.0/8.0 * v(v.size() - 1);
        
            for (int i = 0; i < v.size(); ++i) {
                vFine(2 * i) = v(i);
        
                if (1 < 2 * i + 1 && 2 * i + 1 < nFine - 2) {
                    vFine(2 * i + 1) = 
                        -1.0/16.0 * v(i - 1) + 
                        9.0/16.0 * v(i) + 
                        9.0/16.0 * v(i + 1) - 
                        1.0/16.0 * v(i + 2);
                }
            }
        }    
    }

    

    return vFine;
}

VectorXd MultigridSolver::solveCoarsestGrid(Eigen::SparseMatrix<double> A, VectorXd f) {
    Eigen::SparseLU<Eigen::SparseMatrix<double>> solver;
    solver.compute(A);
    if (solver.info() != Eigen::Success) {
        std::cerr << "SparseLU failed to decompose matrix" << std::endl;
        exit(EXIT_FAILURE);
    }
    return solver.solve(f);
}

VectorXd MultigridSolver::vCycle(Eigen::SparseMatrix<double> A, VectorXd v, VectorXd f, int v1, int v2) {
    /*std::cout << "A.rows() = " << A.rows() << ", A.cols() = " << A.cols() << std::endl;
    std::cout << "v.size() = " << v.size() << ", f.size() = " << f.size() << std::endl;*/

    v = relax(v1, A, v, f);

    if (v.size() <= 4) {
        //std::cout << "Solving on coarsest grid, size = " << v.size() << std::endl;
        v = solveCoarsestGrid(A, f);
    } else {
        VectorXd residual = f - A * v;

        /*std::cout << "Residual size = " << residual.size() << std::endl;*/

        VectorXd fCoarse = restrict(residual);
        VectorXd vCoarse = VectorXd::Zero(fCoarse.size());

        Eigen::SparseMatrix<double> A_coarse(fCoarse.size(), fCoarse.size());
        A_coarse.reserve(Eigen::VectorXi::Constant(fCoarse.size(), 3));

        A_coarse = setupA(fCoarse.size() - 1);

        /*std::cout << "A_coarse.rows() = " << A_coarse.rows() 
                  << ", A_coarse.cols() = " << A_coarse.cols() 
                  << ", vCoarse.size() = " << vCoarse.size() 
                  << ", fCoarse.size() = " << fCoarse.size() << std::endl;*/

        vCoarse = vCycle(A_coarse, vCoarse, fCoarse, v1, v2);
        v += interpolate(vCoarse);
    }

    v = relax(v2, A, v, f);

    /*std::cout << "After relaxation: v.size() = " << v.size() << std::endl;*/

    return v;
}

VectorXd MultigridSolver::fmgCycle(Eigen::SparseMatrix<double> A, VectorXd f, int v1, int v2) {
    int n = f.size();
    VectorXd v;

    // (FMG-1) 如果是最粗网格，直接设置为 0，并进入 V-cycle 修正
    if (n <= 4) {
        v = VectorXd::Zero(n);
    } else {
        // 将源项限制到粗网格
        VectorXd fCoarse = restrict(f);
        VectorXd vCoarse = VectorXd::Zero(fCoarse.size());

        // 创建粗网格的矩阵 A_coarse
        Eigen::SparseMatrix<double> A_coarse(fCoarse.size(), fCoarse.size());
        A_coarse.reserve(Eigen::VectorXi::Constant(fCoarse.size(), 3));
        A_coarse = setupA(fCoarse.size() - 1);  // 构建粗网格矩阵

        // 递归调用 FMG 解决粗网格上的问题
        vCoarse = fmgCycle(A_coarse, fCoarse, v1, v2);

        // (FMG-2) 将粗网格解插值回细网格
        v = interpolate(vCoarse);
    }

    // (FMG-3) 使用 V-cycle 修正
    v = vCycle(A, v, f, v1, v2);

    return v;
}


VectorXd MultigridSolver::solve() {
    VectorXd v;
    VectorXd f;
    Eigen::SparseMatrix<double> A;

    if (dimension == 1) {
        v = Map<VectorXd>(initial_guess.data(), initial_guess.size());
        f = setupVector();
        A = setupMatrix();
    } else if (dimension == 2) {
        auto gridData = grid.getGrid_2D();
        int size = gridData.size();
        v = Map<VectorXd>(initial_guess.data(), size);
        f = setupVector();
        A = setupMatrix();
    }

    if (method == "V-cycle") {
        for (int i = 0; i < maxIterations; ++i) {
            v = vCycle(A, v, f, v1, v2);
            
            // 计算相对误差
            double error = (v - exact_solution).norm() / std::sqrt(v.size());

            // 设置输出格式为科学计数法，并保留4位有效数字
            std::cout << std::scientific << std::setprecision(4);
        
            // 输出当前迭代误差（可选）
            std::cout << "Iteration " << i + 1 << ", Error = " << error << std::endl;
        
            // 如果相对误差小于阈值，提前终止
            if (error < tolerance) {
                std::cout << "Converged after " << i + 1 << " iterations." << std::endl;
                break;
            }
        }
        solution = v;
    } else if (method == "FMG") {
        VectorXd v_0 = VectorXd::Zero(f.size());
        VectorXd f_0 = f;
        for (int i = 0; i < maxIterations; ++i) {
            // 修正右端项
            f = f_0 - A * v_0;
            
            v_0 += fmgCycle(A, f, v1, v2);
            
            // 计算相对误差
            double error = (v_0 - exact_solution).norm() / std::sqrt(v_0.size());

            // 设置输出格式为科学计数法，并保留4位有效数字
            std::cout << std::scientific << std::setprecision(4);
        
            // 输出当前迭代误差
            std::cout << "Iteration " << i + 1 << ", Error = " << error << std::endl;
        
            // 如果相对误差小于阈值，提前终止
            if (error < tolerance) {
                std::cout << "Converged after " << i + 1 << " iterations." << std::endl;
                break;
            }
        }
        solution = v_0;
    }

    return solution;
}

VectorXd MultigridSolver::setupVector() {
    if (dimension == 1) {
        auto g = grid.getGrid_1D();
        int n = g.size();
        int Points = n;
        VectorXd f(Points);

        if (bc.getType() == "Dirichlet"){
            f[0] = bc.getCondition_1D(g[0],0);
            for (int i = 1; i < Points-1; ++i) {
                f[i] = grid.getSourceTerm_1D()(g[i]);
            }
            f[Points-1] = bc.getCondition_1D(g[Points-1],1);
        }

        else if (bc.getType() == "Neumann"){
            f[0] = exact_solution(0);
            for (int i = 1; i < Points-1; ++i) {
                f[i] = grid.getSourceTerm_1D()(g[i]);
            }
            f[Points-1] = bc.getCondition_1D(g[Points-1],1);
        }

        else if(bc.getType() == "mixed"){
            if(types[0] == "Neumann"){
                f[0] = -bc.getCondition_1D(g[0],0);
                for (int i = 1; i < Points-1; ++i) {
                    f[i] = grid.getSourceTerm_1D()(g[i]);
                }
                f[Points-1] = bc.getCondition_1D(g[Points-1],1);
            }
            else if(types[0] == "Dirichlet"){
                f[0] = bc.getCondition_1D(g[0],0);
                for (int i = 1; i < Points-1; ++i) {
                    f[i] = grid.getSourceTerm_1D()(g[i]);
                }
                f[Points-1] = bc.getCondition_1D(g[Points-1],1);
            }
        }

         /*// 添加打印语句
         std::cout << "Vector f dimensions: " << f.size() << std::endl;
         for (int i = 0; i < f.size(); ++i) {
             std::cout << "f[" << i << "] = " << f[i] << std::endl;
         }*/
        return f;
    } else if (dimension == 2) {
        auto g = grid.getGrid_2D();
        int n = g.size() * g[0].size();
        VectorXd f(n);
        // 在这里计算源项 f，取决于具体问题
        return f;
    }
    return VectorXd();
}

Eigen::SparseMatrix<double> MultigridSolver::setupMatrix() {
    if (dimension == 1) {
        auto g = grid.getGrid_1D();
        int n = g.size();
        double h =grid.geth();
        int Points = n;
        Eigen::SparseMatrix<double> A(Points, Points);
        A = setupA(n-1);

        /*// 添加打印语句
        std::cout << "Matrix A dimensions: " << A.rows() << " x " << A.cols() << std::endl;
        for (int i = 0; i < A.outerSize(); ++i) {
            for (Eigen::SparseMatrix<double>::InnerIterator it(A, i); it; ++it) {
                std::cout << "A(" << it.row() << ", " << it.col() << ") = " << it.value() << std::endl;
            }
        
        }*/
        return A;
    } else if (dimension == 2) {
        auto g = grid.getGrid_2D();
        int n = g.size();
        
    }
    Eigen::SparseMatrix<double> emptyMatrix(0, 0);
    return emptyMatrix;
}

Eigen::SparseMatrix<double> MultigridSolver::setupA(int n){
    if (dimension == 1) {
        auto g = grid.getGrid_1D();
        int N = g.size();
        double h =grid.geth();
        int Points = n+1;
        h = N/n * h;
        Eigen::SparseMatrix<double> A(Points, Points);
        A.reserve(Eigen::VectorXi::Constant(Points, 3)); // Reserve space for 3 non-zero entries per row

        if (bc.getType() == "Dirichlet"){
            A.coeffRef(0, 0) = 1.0;
            for (int i = 1; i < n; ++i) {
                A.coeffRef(i, i) = 2.0 / (h * h);
                A.coeffRef(i, i - 1) = -1.0 / (h * h);
                A.coeffRef(i, i + 1) = -1.0 / (h * h);
            }
            A.coeffRef(n, n) = 1.0;
        }

        else if(bc.getType() == "Neumann"){
            A.coeffRef(0, 0) = 1.0;
            for (int i = 1; i < n; ++i) {
                A.coeffRef(i, i) = 2.0 / (h * h);
                A.coeffRef(i, i - 1) = -1.0 / (h * h);
                A.coeffRef(i, i + 1) = -1.0 / (h * h);
            }
            A.coeffRef(n, n-2) = 1.0/(2.0*h);
            A.coeffRef(n, n-1) = -2.0/h;
            A.coeffRef(n, n) = 3.0/(2.0*h);
        }

        else if(bc.getType() == "mixed"){
            if(types[0] == "Neumann"){
                A.coeffRef(0, 0) = -3.0/(2.0*h);
                A.coeffRef(0, 1) = 2.0/h;
                A.coeffRef(0, 2) = -1.0/(2.0*h);
                for (int i = 1; i < n; ++i) {
                    A.coeffRef(i, i) = 2.0 / (h * h);
                    A.coeffRef(i, i - 1) = -1.0 / (h * h);
                    A.coeffRef(i, i + 1) = -1.0 / (h * h);
                }
                A.coeffRef(n, n) = 1.0;
            }
            else if(types[0] == "Dirichlet"){
                A.coeffRef(0, 0) = 1.0;
                for (int i = 1; i < n; ++i) {
                    A.coeffRef(i, i) = 2.0 / (h * h);
                    A.coeffRef(i, i - 1) = -1.0 / (h * h);
                    A.coeffRef(i, i + 1) = -1.0 / (h * h);
                }
                A.coeffRef(n, n-2) = 1.0/(2.0*h);
                A.coeffRef(n, n-1) = -2.0/h;
                A.coeffRef(n, n) = 3.0/(2.0*h);
            }
        }

        return A;
    }
    else if(dimension == 2){
        auto g = grid.getGrid_2D();
        int N = g.size();
        double h =grid.geth();
        int Points = n+1;
        h = N/n * h;
        Eigen::SparseMatrix<double> A(Points, Points);
        A.reserve(Eigen::VectorXi::Constant(Points, 5)); // Reserve space for 5 non-zero entries per row


    }
    Eigen::SparseMatrix<double> emptyMatrix(0, 0);
    return emptyMatrix;
}

double MultigridSolver::compute_error(){
    int n = solution.size();
    double error = (solution - exact_solution).norm() / std::sqrt(n);
    return error;
}

void MultigridSolver::epsilon_sensitivity_test() {
    std::vector<double> epsilons = {1e-4, 1e-6, 1e-8, 1e-10, 1e-12, 1e-14, 2.2e-16};
    
    // 输出表头
    std::cout << "\nε Sensitivity Test\n";
    std::cout << std::setw(12) << "ε" 
              << std::setw(15) << "Error" 
              << std::setw(15) << "Time (s)" 
              << std::setw(15) << "Iterations"
              << std::setw(20) << "Residual Reduction\n";
    std::cout << std::string(77, '-') << std::endl;

    // 保存原始容差
    double original_tolerance = tolerance;
    int original_max_iter = maxIterations;

    for (double eps : epsilons) {
        tolerance = eps;
        maxIterations = 1000;  // 确保有足够迭代次数
        
        // 正确的时间测量
        auto start = std::chrono::steady_clock::now();  // 改用steady_clock
        
        // 重置初始猜测
        solution = VectorXd::Zero(solution.size());
        
        // 求解并捕获迭代次数
        int used_iterations = 0;
        for (; used_iterations < maxIterations; ++used_iterations) {
            solution = vCycle(setupMatrix(), solution, setupVector(), v1, v2);
            double current_error = (setupVector() - setupMatrix() * solution).norm();
            if (current_error < eps) break;
        }
        
        auto end = std::chrono::steady_clock::now();
        std::chrono::duration<double> elapsed = end - start;
        
        double error = compute_error();
        double residual = (setupVector() - setupMatrix() * solution).norm();
        
        // 格式化输出
        std::cout << std::scientific << std::setprecision(3);
        std::cout << std::setw(12) << eps 
                  << std::setw(15) << error 
                  << std::setw(15) << elapsed.count()
                  << std::setw(15) << used_iterations + 1
                  << std::setw(20) << residual/eps << "\n";

        if (error > eps && eps < 1e-10) {  // 只在较小ε时警告
            std::cout << "✖ Failed to achieve tolerance (Error/ε = " 
                      << error/eps << ")\n";
            // 不立即break，继续测试更小的ε
        }
    }

    // 恢复原始设置
    tolerance = original_tolerance;
    maxIterations = original_max_iter;
    
    std::cout << std::defaultfloat;  // 恢复默认输出格式
}
